#  /**
#  * Copyright (c) 2022 Beijing Jiaotong University
#  * PhotLab is licensed under [Open Source License].
#  * You can use this software according to the terms and conditions of the [Open Source License].
#  * You may obtain a copy of [Open Source License] at: [https://open.source.license/]
#  *
#  * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
#  * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
#  * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
#  *
#  * See the [Open Source License] for more details.
#  */
#  /**
#  * Author: Zheng Wang
#  * Created: Jun. 27, 2022
#  * Supported by: National Key Research and Development Program of China
#  */
import logging
import multiprocessing
import time
import psutil
from multiprocessing import Queue
from multiprocessing import Process
from phot_engine.globals.task_status import TaskStatus
from phot_engine.globals.engine_instruction import EngineInstruction
from phot_engine.globals import EngineError
from phot_engine.globals import RunOutOfTime
from phot_engine.logical_structure.resource_manager import ResourceManager
from phot_engine.logical_structure.arithmetic_unit import ArithmeticUnit
from phot_engine.globals.shared_dictionary_key import SharedDictionaryKeyEnum
from phot_engine.globals.arithmetic_unit_status import ArithmeticUnitStatus
from phot_engine.configuration.parse_configuration import load_waiting_strategy_config

waiting_strategy_config_dict = load_waiting_strategy_config()

_MAX_WAITING_TIME = waiting_strategy_config_dict["max_waiting_time"]
_MAX_WAITING_LOOP_NUMBER = waiting_strategy_config_dict["max_waiting_loop_number"]


class Dispatcher(Process):
    """
        Computing engine's dispatch,used to allocate arithmetic unit for vertexes in computational graph
    """
    __process_lock = multiprocessing.Lock()
    __resource_manager = ResourceManager()
    __super_step = 0
    __current_calculate_col_index = 0

    def __init__(self, simulation_task, instruction):
        """

        :param simulation_task: phot_engine.core.simulation_task.SimulationTask type
        :param instruction:
        """
        if simulation_task is not None and instruction is not None:
            super(Dispatcher, self).__init__()
            # initialize communication data structure
            self.__error_queue = Queue()
            # initialize simulation
            self.__simulation_task = simulation_task
            self.__instruction = instruction
            self.__thread_list = []
            # 共享字典
            self.__share_dict = multiprocessing.Manager().dict()
            self.__share_dict[SharedDictionaryKeyEnum.SIMULATION_RESULT_KEY] = []
            self.recommend_gpu = self.__get_recommendGpu()
    def get_error_queue(self):
        return self.__error_queue

    def get_task(self):
        """
            获取调度器中的仿真任务
            get simulation task in dispatcher
        :return:
        """
        return self.__simulation_task

    def get_task_status(self):
        return self.__simulation_task.get_task_status()

    def alter_instruction(self, new_instruction):
        """
            修改计算引擎计算指令
            alter engine instruction in the dispatcher
        :param new_instruction:
        :return:
        """
        self.__instruction = new_instruction

    def __detect_exception(self):
        """
            detecting if current dispatcher has exceptions
        :return:
        """

        if not self.__error_queue.empty():
            return True
        else:
            return False

    def execute_instruction(self):
        """
            执行计算引擎指令
            Execute computing engine instruction
        :return:
        """
        try:
            if self.__instruction == EngineInstruction.EXECUTE:
                logging.info("According to engine instruction-Prepare to execute the simulation task {}."
                             .format(self.__simulation_task.get_task_id()))

                self.__simulation_task.set_task_status(TaskStatus.TASK_RUNNING)
                self.start()
                self.join()
                if self.__detect_exception():
                    self.__simulation_task.set_task_status(TaskStatus.TASK_DIED)
                else:
                    # 正常执行设置结束时间与任务完成状态
                    self.__simulation_task.set_end_executing_time(time.time())
                    self.__simulation_task.set_task_status(TaskStatus.TASK_ACCOMPLISHED)
                    self.__simulation_task.set_edge_list(
                        self.__share_dict[SharedDictionaryKeyEnum.SIMULATION_RESULT_KEY])
            elif self.__instruction == EngineInstruction.HALT:
                logging.info("According to engine instruction-Prepare to pause the simulation task {}."
                             .format(self.__simulation_task.get_task_id()))

                self.__simulation_task.set_task_status(TaskStatus.TASK_STOPPED)
                pause = psutil.Process(self.pid)
                pause.suspend()
            elif self.__instruction == EngineInstruction.RECOVER:
                logging.info("According to engine instruction-Prepare to recover the simulation task {}."
                             .format(self.__simulation_task.get_task_id()))
                self.__simulation_task.set_task_status(TaskStatus.TASK_RUNNING)
                pause = psutil.Process(self.pid)
                pause.resume()
            elif self.__instruction == EngineInstruction.CANCEL:
                self.__super_step = 0
                self.__current_calculate_col_index = 0
                self.__simulation_task.clear_all_signal()
        except Exception as e:
            pass

    def __calculate_task(self):
        """
            计算任务
            calculate task
        :return:
        """
        try:
            graph = self.__simulation_task.get_graph()
            computational_graph = self.__simulation_task.get_computational_graph()
            global_parameter = self.__simulation_task.get_global_parameter()
            edge_list = self.__simulation_task.get_edge_list()
            max_row, max_col = computational_graph.shape
            self.__super_step = 0
            start_time = time.time()
            while self.__current_calculate_col_index < max_col:
                """
                按列遍历计算图,当全部列都遍历结束后，完成系统计算
                The calculation chart is traversed by column. 
                When all columns are traversed, the system calculation is completed
                """
                running_thread_number = self.__count_running_thread_number()
                if 0 == running_thread_number:
                    # 上一批线程全部计算结束,为下一批图顶点分配运算单元
                    self.__allocate_arithmetic_unit(graph=graph,
                                                    global_parameter=global_parameter,
                                                    computational_graph=computational_graph,
                                                    edge_list=edge_list,
                                                    max_row=max_row)
                    self.__super_step += 1    #等待次数
                    self.__current_calculate_col_index += 1
                else:
                    # 若有一些线程还未完成计算
                    self.__wait_calculating_thread()

            end_time = time.time()
            self.__simulation_task.set_end_executing_time(time.time())
            logging.info("Simulation task has finished calculation.Synchronous time is:{},time usage is:{}"
                         .format(self.__super_step, (end_time - start_time)))
            return self.__simulation_task.get_edge_list()
        except RunOutOfTime as run_out_of_time:
            logging.error("Computing engine run out of time!!!")
            self.__simulation_task.set_task_status(TaskStatus.TASK_DIED)
            self.close()
            raise run_out_of_time
        except Exception as e:
            logging.error("An exception occurs in computing engine :{}".format(e))
            self.__simulation_task.set_task_status(TaskStatus.TASK_DIED)
            self.close()
            raise e
        finally:
            pass

    def __count_running_thread_number(self) -> int:
        """
        检查线程列表中的线程是否存活,检查进程计算状况,返回存活线程和未结束线程数量
        Check whether the thread in the thread list is alive.Return alive thread number
        :return:
        """
        if len(self.__thread_list) == 0:
            return 0
        else:
            alive_thread_cnt = 0
            for one_thread in self.__thread_list:
                # print("线程列表扫描时间:{},线程名:{},运行顶点:{},是否存活:{},线程运算状态:{}".format(time.time(),
                #                                                                     one_thread.getName(),
                #                                                                     one_thread.get_vertex_id(),
                #                                                                     one_thread.is_alive(),
                #                                                                     one_thread.get_arithmetic_unit_status()))
                if one_thread.get_arithmetic_unit_status() == ArithmeticUnitStatus.THREAD_RUNNING \
                        or one_thread.get_arithmetic_unit_status() == ArithmeticUnitStatus.THREAD_INITIALIZED:
                    alive_thread_cnt += 1
                elif one_thread.get_arithmetic_unit_status() == ArithmeticUnitStatus.THREAD_EXCEPTION:
                    raise EngineError("Arithmetic Unit:{},Vertex Id:{} has Exception"
                                      .format(one_thread.getName(), one_thread.get_vertex_id()))
                elif one_thread.get_arithmetic_unit_status() == ArithmeticUnitStatus.THREAD_ACCOMPLISHED:
                    self.__thread_list.remove(one_thread)
                    return len(self.__thread_list)
                elif not one_thread.is_alive():
                    self.__thread_list.remove(one_thread)
                    return len(self.__thread_list)

            return alive_thread_cnt

    def __wait_calculating_thread(self):
        """
            获取 self.__thread_list,根据线程存活数量调整同步等待时间
        :return:
        """
        logging.debug("There are some threads (number:{}) hasn't finished calculation.Waiting time is:{}"
                      .format(len(self.__thread_list), len(self.__thread_list) * _MAX_WAITING_TIME))

        loop_cnt = 1
        waiting_thread_number = self.__count_running_thread_number()
        while waiting_thread_number > 0 and loop_cnt <= _MAX_WAITING_LOOP_NUMBER:
            waiting_time = _MAX_WAITING_TIME * waiting_thread_number
            for one_arithmetic_thread in self.__thread_list:
                logging.debug("Arithmetic Unit thread:{} is calculating vertex:{},waiting time is {} s"
                              .format(one_arithmetic_thread.getName(),
                                      one_arithmetic_thread.get_vertex_id(),
                                      waiting_time))

                one_arithmetic_thread.join(waiting_time)
                loop_cnt += 1
                one_arithmetic_thread.get_vertex().set_super_step(self.__super_step)

            waiting_thread_number = self.__count_running_thread_number()
            loop_cnt += 1
        if loop_cnt > _MAX_WAITING_LOOP_NUMBER:
            raise RunOutOfTime("Run out of time")
    def __get_recommendGpu(self):
         return self.__resource_manager.get_gpu_usage()

    def __allocate_arithmetic_unit(self, graph, global_parameter, computational_graph, edge_list, max_row):
        for row_index in range(0, max_row):
            # 按行对图顶点开若干线程进行并发计算
            if computational_graph[row_index][self.__current_calculate_col_index] is not None:
                one_vertex = graph.get_vertex(
                    computational_graph[row_index][self.__current_calculate_col_index])

                # New an arithmetic thread
                # 目前是任务级
                arithmetic_thread = ArithmeticUnit(vertex=one_vertex,
                                                   edge_list=edge_list,
                                                   error_queue=self.__error_queue,
                                                   recommend_gpu=self.recommend_gpu)

                arithmetic_thread.set_global_parameter(global_parameter)
                logging.debug("Allocate an arithmetic unit({}) thread for vertex:{} in super step:{}"
                              .format(arithmetic_thread.getName(),
                                      arithmetic_thread.get_vertex_id(),
                                      self.__super_step))

                self.__thread_list.append(arithmetic_thread)
                arithmetic_thread.start()

    def set_worker_cpu(self):
        """
            set a list of worker cpu to this dispatcher
        :return:
        """
        self.__process_lock.acquire()
        worker_cpu = self.__resource_manager.recommend_cpus()
        process_priority = self.__resource_manager.recommend_priority(self.__simulation_task.get_priority())
        one_process = psutil.Process(self.pid)
        one_process.cpu_affinity(worker_cpu)
        one_process.nice(value=process_priority)
        self.__process_lock.release()
        logging.debug("Current process pid is {},CPU affinity is {}".format(self.pid, one_process.cpu_affinity()))



    def run(self) -> None:
        try:
            logging.info("Start to calculate simulation task:{}.Allocated dispatcher is {}"
                         .format(self.__simulation_task.get_task_id(), self.pid))
            self.set_worker_cpu()
            self.__share_dict[SharedDictionaryKeyEnum.SIMULATION_RESULT_KEY] = self.__calculate_task()

        except RunOutOfTime as run_out_of_time:
            self.__error_queue.put(run_out_of_time)
            raise run_out_of_time
        except Exception as e:
            self.__error_queue.put(e)
            raise e
        finally:
            self.__simulation_task.set_task_status(TaskStatus.TASK_DIED)
            # # 信号0 探测进程存在性
            # os.kill(self.pid, 0)
